import math

from dl.DNN.dataset import DNN_Dataset
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
# ['trade_date', 'close', 'open', 'high', 'low', 'pre_close', 'change',
#        'pct_chg', 'vol', 'amount', 'turnover_rate','MACD', 'DEA', 'Histogram', 'MA5', 'MA10',
#        'RSI', 'MOM', 'EMA12', 'EMA26', 'open-close', 'high-low', 'Fmark']
def sign(x):
    if x > 0:return 1
    return -1

"""
训练与预测DNN网络
"""
class DNN_sk(object):
    def __init__(self, ts_code, start_date, end_date):
        self.tmpdata = DNN_Dataset(ts_code = ts_code, start_date=start_date, end_date=end_date)
        self.data = self.tmpdata.data
        self.label_cols = self.data.columns.drop('trade_date')  # 用于训练模型的特征
        self.model = MLPClassifier(solver='lbfgs', alpha=1e-5, max_iter=500, hidden_layer_sizes=3 * [500], random_state=1)
        self.traindata, self.testdata = train_test_split(self.data, test_size = 0.5, random_state=100)
        #print('count---', (self.traindata['direction'] == 1).value_counts())
        self.traindata, self.testdata = self.traindata.copy().sort_index(), self.testdata.copy().sort_index()

    def train(self):
        self.model.fit(self.traindata[self.label_cols], self.traindata['direction'])

    def predict(self):
        self.testdata['pos_dnn_sk'] = self.model.predict(self.testdata[self.label_cols]) # pos_dnn_sk是预测出来的值
        self.testdata['start_dnn_sk'] = self.testdata['pos_dnn_sk'] * self.testdata['returns'] # 将预测值乘上returns表示收益
        self.testdata[['start_dnn_sk', 'returns']].cumsum().plot(figsize = (10, 6)) # 将收益做前缀和，画出图形，图像越高，预测结果越好
        plt.show()

if __name__ == '__main__':
    dnn = DNN_sk('000001.SZ', '20200102', '20210102')
    dnn.train()
    dnn.predict()